# https://docs.langchain.com/oss/python/langgraph/overview
# pip install -U langgraph

# NOTE:
# 1. 上面overview原谅中，MessageGraph在1.0.0中已被废弃，

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# TypedDict：用于定义强类型的状态结构（State类），明确状态字段的类型约束。
class State(TypedDict):
    # Annotated + add_messages：LangGraph 的消息处理注解，用于自动管理messages列表的更新（如追加消息，而非覆盖）。
    # 字段类型为list，用于存储对话历史（如用户消息、AI 响应）。
    # add_messages注解是关键：告知 LangGraph，当节点返回{"messages": [...]}时，
    #    自动将新消息追加到现有messages列表中，而非直接替换（避免丢失历史对话）。
    messages: Annotated[list, add_messages]

# 定义图的核心业务节点（模拟 LLM 的响应生成逻辑），节点是图的基本执行单元
# 模拟 LLM 处理用户消息，此处简化为固定返回 "Hello, World!"。
def mock_llm(state: State):
    print("节点：mock_llm")
    return {"messages": [{"role": "ai", "content":"我是大语言模型"}]}

def main():
    # 1.构建图
    # StateGraph：图构建器，用于定义节点、边和状态流转规则。
    graph = StateGraph(State)
    graph.add_node("mock_llm", mock_llm)
    graph.add_edge(START, "mock_llm")
    graph.add_edge("mock_llm", END)
    # 编译过程会验证节点、边的合法性（如节点输出是否符合State结构）。
    # 生成优化后的执行逻辑，支持invoke（同步调用）、astream（异步流式调用）等方法。
    graph = graph.compile()

    # 2.执行图并返回结果
    # 触发图的同步执行，传入初始状态
    # 返回的result类型是 LangChain 框架中的 HumanMessage 和 AIMessage 对象
    result = graph.invoke({"messages":[{"role": "user", "content": "Hi"}]})

    # 3.打印结果
    print("\n===== 完整状态结果 =====")
    print(result)  # 打印整个最终状态（包含messages字段）

    print("\n===== 所有对话消息 =====")
    for msg in result["messages"]:
        # HumanMessage 或 AIMessage 实例（对象），而非普通字典。
        # 这类对象的属性（如 content）需要通过 . 符号 访问，而非字典的 [] 下标。
        print(f"{msg.type}: {msg.content}")


    print("\n===== AI的回复 =====")
    # 提取最后一条消息（AI的回复）
    # HumanMessage/AIMessage 是对象，不是字典，必须用 .属性名 访问内容（如 .content）和角色（如 .role）。
    ai_msg = result["messages"][-1].content
    print(ai_msg)

if __name__ == "__main__":
    main()
